import cv2

# TO CHANGE！
img_file = "image/203.jpg"


points = []
scale_x = scale_y = 1.0  # 缩放比例


def mouse_callback(event, x, y, flags, param):
    global points, img_display, scale_x, scale_y

    if event == cv2.EVENT_LBUTTONDOWN:
        # 映射回原图坐标
        orig_x = int(x / scale_x)
        orig_y = int(y / scale_y)
        points.append((orig_x, orig_y))
        print(f"点 {len(points)}: ({orig_x}, {orig_y})")

        # 在显示图上画（画的点还是显示图的坐标）
        cv2.circle(img_display, (x, y), 5, (0, 255, 0), -1)
        if len(points) > 1:
            prev_x = int(points[-2][0] * scale_x)
            prev_y = int(points[-2][1] * scale_y)
            cv2.line(img_display, (prev_x, prev_y), (x, y), (0, 0, 255), 2)

        cv2.imshow("image", img_display)

def main():
    global img_display, points, scale_x, scale_y
    name = img_file.split("/")[-1].split("\\")[-1].split(".")[-2]
    img = cv2.imread(img_file)  # 原始图
    if img is None:
        print("读取图片失败")
        return

    # 缩放到 600×800
    target_w, target_h = 800, 450
    scale_x = target_w / img.shape[1]
    scale_y = target_h / img.shape[0]
    img_display = cv2.resize(img, (target_w, target_h))

    cv2.namedWindow("image")
    cv2.setMouseCallback("image", mouse_callback)

    cv2.imshow("image", img_display)

    while True:
        key = cv2.waitKey(1) & 0xFF
        if key == 13:  # Enter
            if len(points) > 2:
                # 显示闭合效果（要映射到显示坐标）
                first_x = int(points[0][0] * scale_x)
                first_y = int(points[0][1] * scale_y)
                last_x = int(points[-1][0] * scale_x)
                last_y = int(points[-1][1] * scale_y)
                cv2.line(img_display, (last_x, last_y), (first_x, first_y), (0, 0, 255), 2)
                cv2.imshow("image", img_display)
                cv2.waitKey(500)
            break
        elif key == ord('q'):
            break

    cv2.destroyAllWindows()

    # 保存原图坐标
    with open(f"{name}_location.txt", "w") as f:
        for (x, y) in points:
            f.write(f"[{x},{y}],")
    print(f"点坐标已保存到 {name}_location.txt")

if __name__ == "__main__":
    main()
